import pandas as pd
import numpy as np
import statsmodels.api as sm
import matplotlib
# 设置matplotlib使用非交互式后端，避免tkinter相关错误
matplotlib.use('Agg')
import matplotlib.pyplot as plt

class PriceElasticityAnalyzer:
    def __init__(self, data=None):
        self.data = data
        self.elasticity = None
        self.model = None
        
    def calculate_elasticity(self, data=None, price_col='avg_price', quantity_col='sales_quantity'):
        """计算价格弹性"""
        if data is not None:
            self.data = data
            
        # 确保数据是DataFrame类型
        if not isinstance(self.data, pd.DataFrame):
            raise ValueError("数据必须是pandas DataFrame类型")
            
        # 取对数转换，为计算弹性做准备
        log_data = self.data.copy()
        log_data['log_price'] = np.log(self.data[price_col])
        log_data['log_quantity'] = np.log(self.data[quantity_col] + 1)  # 加1避免log(0)
        
        # 删除缺失值
        log_data = log_data.dropna(subset=['log_price', 'log_quantity'])
        
        # 构建OLS模型: log(quantity) = b0 + b1*log(price)
        X = log_data['log_price']
        X = sm.add_constant(X)
        y = log_data['log_quantity']
        
        # 拟合模型
        self.model = sm.OLS(y, X).fit()
        
        # 弹性系数就是价格的系数(斜率)
        self.elasticity = self.model.params['log_price']
        
        return {
            'elasticity': self.elasticity,
            'model_summary': self.model.summary()
        }
    
    def plot_elasticity(self, save_path=None):
        """绘制价格与销量的关系图"""
        if self.data is None:
            raise ValueError("必须先加载数据")
            
        price_col = 'avg_price'
        quantity_col = 'sales_quantity'
        
        fig, ax = plt.subplots(figsize=(10, 6))
        
        # 散点图
        ax.scatter(self.data[price_col], self.data[quantity_col], alpha=0.5)
        
        # 如果已经计算了弹性，添加拟合线
        if self.model is not None:
            # 生成预测用的价格范围
            price_range = np.linspace(self.data[price_col].min(), self.data[price_col].max(), 100)
            log_price_range = np.log(price_range)
            
            # 添加常数项
            X_pred = sm.add_constant(log_price_range)
            
            # 预测对数销量
            log_quantity_pred = self.model.predict(X_pred)
            
            # 转换回原始尺度
            quantity_pred = np.exp(log_quantity_pred) - 1  # 减1是因为之前加1
            
            # 绘制拟合线
            ax.plot(price_range, quantity_pred, 'r-', linewidth=2)
            
            # 添加弹性系数标注
            elasticity_text = f'价格弹性系数: {self.elasticity:.2f}'
            interpretation = '富有弹性' if abs(self.elasticity) > 1 else '缺乏弹性'
            ax.text(0.05, 0.95, f'{elasticity_text}\n({interpretation})', 
                   transform=ax.transAxes, fontsize=12, verticalalignment='top')
        
        ax.set_title('价格与销量的关系')
        ax.set_xlabel('价格')
        ax.set_ylabel('销量')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path)
        
        # 关闭图形，避免内存泄漏和tkinter错误
        result_fig = fig
        plt.close(fig)
            
        return result_fig
    
    def analyze_by_category(self):
        """按类别分析价格弹性"""
        if self.data is None:
            raise ValueError("必须先加载数据")
            
        # 确保数据包含类别信息
        if 'category' not in self.data.columns:
            raise ValueError("数据必须包含'category'列")
            
        # 按类别分组并计算弹性
        categories = self.data['category'].unique()
        results = {}
        
        for category in categories:
            cat_data = self.data[self.data['category'] == category]
            
            # 如果数据点太少，跳过
            if len(cat_data) < 5:
                continue
                
            # 计算弹性
            try:
                elasticity_result = self.calculate_elasticity(cat_data)
                results[category] = elasticity_result['elasticity']
            except Exception as e:
                print(f"计算类别 {category} 的弹性时出错: {str(e)}")
        
        return results
